package main

import (
	"flag"
	"fmt"
	"gopkg.in/yaml.v3"
	"gorm.io/driver/mysql"
	"gorm.io/driver/postgres"
	"gorm.io/driver/sqlite"
	"gorm.io/driver/sqlserver"
	"gorm.io/gen"
	"gorm.io/gorm"
	"log"
	"os"
	"sync"
)

// DBType database type
type DBType string

const (
	// dbMySQL Gorm Drivers mysql || postgres || sqlite || sqlserver
	dbMySQL     DBType = "mysql"
	dbPostgres  DBType = "postgres"
	dbSQLite    DBType = "sqlite"
	dbSQLServer DBType = "sqlserver"
)
const (
	// DefaultOutPath default path
	DefaultOutPath = "./dao/query"
)

// CmdParams is command line parameters
type CmdParams struct {
	DSN               string   `yaml:"dsn"`               // consult[https://gorm.io/docs/connecting_to_the_database.html]"
	DB                string   `yaml:"db"`                // input mysql or postgres or sqlite or sqlserver. consult[https://gorm.io/docs/connecting_to_the_database.html]
	Tables            []string `yaml:"tables"`            // enter the required data table or leave it blank
	OnlyModel         bool     `yaml:"onlyModel"`         // only generate model
	OutPath           string   `yaml:"outPath"`           // specify a directory for output
	OutFile           string   `yaml:"outFile"`           // query code file name, default: gen.go
	WithUnitTest      bool     `yaml:"withUnitTest"`      // generate unit test for query code
	ModelPkgName      string   `yaml:"modelPkgName"`      // generated model code's package name
	FieldNullable     bool     `yaml:"fieldNullable"`     // generate with pointer when field is nullable
	FieldWithIndexTag bool     `yaml:"fieldWithIndexTag"` // generate field with gorm index tag
	FieldWithTypeTag  bool     `yaml:"fieldWithTypeTag"`  // generate field with gorm column type tag
	FieldSignable     bool     `yaml:"fieldSignable"`     // detect integer field's unsigned type, adjust generated data type
}

// YamlConfig is yaml config struct
type YamlConfig struct {
	Version  string                `yaml:"version"`  //
	Database *map[string]CmdParams `yaml:"database"` //
}

// connectDB choose db type for connection to database
func connectDB(t DBType, dsn string) (*gorm.DB, error) {
	if dsn == "" {
		return nil, fmt.Errorf("dsn cannot be empty")
	}

	switch t {
	case dbMySQL:
		return gorm.Open(mysql.Open(dsn))
	case dbPostgres:
		return gorm.Open(postgres.Open(dsn))
	case dbSQLite:
		return gorm.Open(sqlite.Open(dsn))
	case dbSQLServer:
		return gorm.Open(sqlserver.Open(dsn))
	default:
		return nil, fmt.Errorf("unknow db %q (support mysql || postgres || sqlite || sqlserver for now)", t)
	}
}

// genModels is gorm/gen generated models
func genModels(g *gen.Generator, db *gorm.DB, tables []string) (models []interface{}, err error) {
	var tablesList []string
	if len(tables) == 0 {
		// Execute tasks for all tables in the database
		tablesList, err = db.Migrator().GetTables()
		if err != nil {
			return nil, fmt.Errorf("GORM migrator get all tables fail: %w", err)
		}
	} else {
		tablesList = tables
	}

	// Execute some data table tasks
	models = make([]interface{}, len(tablesList))
	for i, tableName := range tablesList {
		models[i] = g.GenerateModel(tableName)
	}
	return models, nil
}

// loadConfigFile load config file from path
func loadConfigFile(path string) (*map[string]CmdParams, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer file.Close() // nolint
	var yamlConfig YamlConfig
	if cmdErr := yaml.NewDecoder(file).Decode(&yamlConfig); cmdErr != nil {
		return nil, cmdErr
	}
	return yamlConfig.Database, nil
}

// argParse is parser for cmd
func argParse() (configFileParams *map[string]CmdParams) {
	// choose is file or flag
	genPath := flag.String("f", "gen.yml", "is path for gen.yml")
	flag.Parse()
	if !Exist(*genPath) {
		log.Printf("Cannot find the specified file “%s”\n", *genPath)
		return
	}

	var err error
	if configFileParams, err = loadConfigFile(*genPath); err == nil && configFileParams != nil {
		log.Printf("%+v\n", configFileParams)
	}

	return configFileParams
}

func execute(conf CmdParams) {
	if conf.OutPath == "" {
		conf.OutPath = DefaultOutPath
	}
	db, err := connectDB(DBType(conf.DB), conf.DSN)
	if err != nil {
		log.Fatalln("connect db server fail:", err)
	}

	g := gen.NewGenerator(gen.Config{
		OutPath:           conf.OutPath,
		OutFile:           conf.OutFile,
		ModelPkgPath:      conf.ModelPkgName,
		WithUnitTest:      conf.WithUnitTest,
		FieldNullable:     conf.FieldNullable,
		FieldWithIndexTag: conf.FieldWithIndexTag,
		FieldWithTypeTag:  conf.FieldWithTypeTag,
		FieldSignable:     conf.FieldSignable,
		Mode:              gen.WithDefaultQuery | gen.WithQueryInterface,
	})

	g.UseDB(db)

	models, err := genModels(g, db, conf.Tables)
	if err != nil {
		log.Fatalln("get tables info fail:", err)
	}

	if !conf.OnlyModel {
		g.ApplyBasic(models...)
	}
	g.Execute()
	wg.Done()
}

var wg sync.WaitGroup

// Exist 文件或文件夹是否存在
func Exist(path string) bool {
	_, err := os.Stat(path)
	if err == nil {
		return true
	}
	return os.IsExist(err)
}

func main() {
	// cmdParse
	config := argParse()
	if config == nil {
		log.Fatalln("parse config fail")
	}
	for _, conf := range *config {
		wg.Add(1)
		go execute(conf)
	}
	wg.Wait()
}
